from typing import Union
from lwpackage.lwdata import get_risk_free_rate
from .factor import *


def get_factor_values(Data: pd.DataFrame,
                      fc_number_list: Union[int, list]):
    """
    return a dataframe will factor values in fc_number_list
    :param Data:
    :param fc_number_list:
    :return:
    """

    assert Data.index.names[0] == 'time'
    assert Data.index.names[1] == 'instrument_id'
    if isinstance(fc_number_list, int):
        fc_number_list = [fc_number_list]
    df = Data.copy()
    fc_name_list = []
    for fc_number in fc_number_list:
        fc_name, mapper = _get_one_factor_value(df, fc_number)
        if fc_name in df.columns:
            df.drop(columns=fc_name, inplace=True)
        df[fc_name] = df.index.map(mapper)
        fc_name_list.append(fc_name)

    return df, fc_name_list


def _get_one_factor_value(Data, fc_number):
    df = Data.copy()
    mapper = (
        df
        .groupby('instrument_id', as_index=False, sort=False)
        .apply(eval(f'factormaker{fc_number}'))
        .reset_index().set_index(['time', 'instrument_id']).iloc[:, -1]
    )
    if mapper.name == 'factor':
        fc_name = mapper.name + fc_number
    elif mapper.name == 'fc':
        fc_name = 'factor' + fc_number
    else:
        fc_name = mapper.name

    return fc_name, mapper


def get_ret(Data: pd.DataFrame,
            fc_freq: str,
            transaction_period: Union[int, None] = None,
            ret_freq: int = 30,
            rfr: bool = True):
    """
    Calculate return based on average transaction price.

    :param Data: panel price Data with open, high, low, close price
    :param fc_freq: fc_freq
    :param transaction_period: transaction_period
    :param ret_freq: ret_freq
    :param rfr: considering risk free rate or not
    :return: a dataframe with a ret column
    """
    df = Data.copy()
    assert df.index.names[0] == 'time'
    assert df.index.names[1] == 'instrument_id'

    # special logic for transaction period
    if not transaction_period:
        if fc_freq == '1d':
            transaction_period = 1
        elif fc_freq == '5m':
            transaction_period = 1
        else:  # 1m freq
            transaction_period = 3

    assert ret_freq > 0
    assert transaction_period > 0
    assert ret_freq >= 2 * transaction_period, f'Transaction period is ' \
                                               f'{transaction_period}, ' \
                                               f'while ret freq is {ret_freq}, invalid!'

    # transaction price of one bar is the average transaction price for a transaction that is completed at
    # the close time of this bar.
    df['transaction_price'] = \
        df[['open', 'high', 'low', 'close']].mean(axis=1).rolling(
            transaction_period).mean()

    mapper = df.groupby('instrument_id')['transaction_price'].apply(
        lambda x: (x.shift(-ret_freq - 1) / x.shift(-transaction_period + 1) - 1)).droplevel(0)

    df['ret'] = df.index.map(mapper)
    df = df.sort_values(by='time')

    if rfr:  # risk-free rate will have effect on sharpe only if long_only == True
        # annualized risk-free rate
        # we fill nan values of risk-free rate with next day's risk-free rate if available
        df.reset_index(inplace=True)
        df['date'] = df['time'].dt.date

        df_rfr = get_risk_free_rate()[['date', 'rate']].copy()
        df_rfr['date'] = df_rfr['date'].dt.date
        df = df.merge(df_rfr, on='date', how='left', validate='m:1')
        df['rate'] = df['rate'].fillna(method='bfill').fillna(method='ffill')
        df.rename(columns={'rate': 'rfr'}, inplace=True)
        df.set_index(['time', 'instrument_id'], inplace=True)

        # simple compound
        df['rfr'] = df['rfr'] / 252

        # excess ret = ret - risk free rate
        df['ret'] = df['ret'] - df['rfr']
        df.drop(columns=['rfr', 'date'], inplace=True)
        df.dropna(how='any', axis=0, inplace=True)

    return df
